package cn.huangsm.advance.ratelimiter2.aop;

import cn.huangsm.advance.ratelimiter2.AccessRateLimiter;
import cn.huangsm.advance.ratelimiter2.utils.Result;
import com.alibaba.fastjson.JSON;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Scope;
import org.springframework.stereotype.Component;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.Map;

/**
 * @author huang
 * @PACKAGE_NAME cn.huangsm.advance.ratelimiter2.aop
 * @PROJECT_NAME advance-code
 * @date 2019/1/19
 */
@Aspect
@Component
@Scope
@Slf4j
public class RateLimiter2Aspect {

    /**
     * 用来存放不同接口的RateLimiter(key为接口名称，value为RateLimiter)
     */
    Map<String, RateLimiter> map = Maps.newConcurrentMap();
    private RateLimiter rateLimiter;

    @Autowired
    private HttpServletResponse response;

    @Pointcut("@annotation(cn.huangsm.advance.ratelimiter2.AccessRateLimiter)")
    public void rateLimiterPointCut() {
    }

    ;

    @Around("rateLimiterPointCut()")
    public Object aroundRateLimiter(ProceedingJoinPoint joinPoint) {
        Object obj = null;
        //获取拦截的方法名
        Signature signature = joinPoint.getSignature();
        //获取拦截的方法名
        MethodSignature methodSignature = (MethodSignature) signature;

        //获取注解类
        AccessRateLimiter annotation = methodSignature.getMethod().getAnnotation(AccessRateLimiter.class);
        //获取注解每秒加入桶中的token
        double limitNum = annotation.limitNum();
        //注解所在方法名区分不同的限流策略
        String functionName = methodSignature.getName();
        //获取rateLimiter
        if (map.containsKey(functionName)) {
            rateLimiter = map.get(functionName);
        } else {
            map.put(functionName, RateLimiter.create(limitNum));
            rateLimiter = map.get(functionName);
        }

        try {
            if (rateLimiter.tryAcquire()) {
                //执行方法
                obj = joinPoint.proceed();
            } else {
                //拒绝了请求（服务降级）
                String result = JSON.toJSONString(Result.ERROR(500, "系统繁忙！"));
                log.info("拒绝了请求：" + result);
                outErrorResult(result);
            }
        } catch (Throwable throwable) {
            throwable.printStackTrace();
        }
        return obj;
    }

    /**
     * 将结果返回
     *
     * @param result
     */
    public void outErrorResult(String result) {
        response.setContentType("application/json;charset=UTF-8");
        try (ServletOutputStream outputStream = response.getOutputStream()) {
            outputStream.write(result.getBytes("utf-8"));
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}
